#pragma once
#include <iostream>
#include <memory>
#include <string>

enum Color { RED, BLACK };

template<class K, class V>
struct Node {
	std::pair<K, V> data;
	Color color;
	std::shared_ptr<Node> left, right;
	std::weak_ptr<Node> parent;

	Node(std::pair<K,V> data) : data(data), color(RED) {}
};


template<class K, class V>
class RedBlackTree {
private:
	std::shared_ptr<Node<K,V>> root;

	void leftRotate(std::shared_ptr<Node<K,V>> x);
	void rightRotate(std::shared_ptr<Node<K,V>> y);
	void insertFixup(std::shared_ptr<Node<K,V>> z);
	void deleteFixup(std::shared_ptr<Node<K,V>> x);
	void transplant(std::shared_ptr<Node<K,V>> u, std::shared_ptr<Node<K,V>> v);
	std::shared_ptr<Node<K,V>> minimum(std::shared_ptr<Node<K,V>> x);

public:
	RedBlackTree() : root(nullptr) {}
	void insert(K key, V value);
	void remove(K key);
	void modify(K key, V new_value);
	void inorder();
	void inorderHelper(std::shared_ptr<Node<K,V>> node);

	class iterator {
	private:
		std::shared_ptr<Node<K,V>> current;
		std::shared_ptr<Node<K,V>> root;

		void increment() {
			if (current == nullptr) return;
			if (current->right) {
				current = current->right;
				while (current->left) current = current->left;
			} else {
				auto parent = current->parent.lock();
				while (parent && current == parent->right) {
					current = parent;
					parent = current->parent.lock();
				}
				current = parent;
			}
		}

	public:
		iterator(std::shared_ptr<Node<K,V>> node, std::shared_ptr<Node<K,V>> root)
		: current(node), root(root) {}

		std::pair<const K&, V&> operator*() const {
			return std::pair<const K&, V&>(current->data.first, current->data.second);
		}

		iterator& operator++() {
			increment();
			return *this;
		}

		iterator operator++(int) {
			iterator tmp = *this;
			increment();
			return tmp;
		}

		bool operator==(const iterator& other) const {
			return current == other.current;
		}

		bool operator!=(const iterator& other) const {
			return !(*this == other);
		}
	};

	iterator begin() {
		if (!root) return end();
		auto node = root;
		while (node->left) node = node->left;
		return iterator(node, root);
	}

	iterator end() {
		return iterator(nullptr, root);
	}
};

RedBlackTree<std::string, std::string> newStrRBTree() {
    return RedBlackTree<std::string, std::string>();
}